import os
import pickle
import numpy as np

import lmdb
import caffe

N = 16
M = 64

TRAIN_POS_NUM = 800
TRAIN_NEG_NUM = 60000

name = "subbands"

def delete_lmdb():
    # os.system("rm -r /home/sunzy/workspace/data/MedlatTrainingData/lmdb/subints_train_lmdb")
    os.system("rm -r /usr/data/MedlatTrainingData/lmdb/%s_train_lmdb" % name)

def load_into_lmdb():
    delete_lmdb()

    # root_path = '/home/sunzy/workspace/data/MedlatTrainingData/feature'
    root_path = '/usr/data/MedlatTrainingData/feature'
    # load pulsars feature from npy file
    with open(os.path.join(root_path, "pulsars_%s") % name, 'r') as f:
        pulsar_feature = pickle.load(f)
    print(len(pulsar_feature))
    for k, v in pulsar_feature.items():
        # print(v.shape)
        pulsar_feature[k] = v[:N, :M].reshape((1, N, M))

    # split data sets into two parts: train and test
    pos_index = pulsar_feature.keys()
    # with open('/home/sunzy/workspace/data/MedlatTrainingData/map/pos_index.txt', 'w') as f:
    with open('/usr/data/MedlatTrainingData/map/pos_index.txt', 'w') as f:
        for i in range(len(pos_index)):
            f.write("%d, %s\n" % (i, pos_index[i]))
    train_pos_index = pos_index[:TRAIN_POS_NUM]
    test_pos_index = pos_index[TRAIN_POS_NUM:]
    # print(len(train_pos_index), len(test_pos_index))
    # for id in train_pos_index:
    #     print(id, pulsar_feature[id].shape)

    # load nonpulsar feature
    nonpulsar_feature = {}
    for i in range(1):# 9 for windows
        dir_path = "RFI_%s" % name
        file_path = os.path.join(root_path, dir_path)
        # for py3
        # with open(file_path, 'rb') as f:
        # for py2
        with open(file_path, 'rb') as f:
            temp_dict = pickle.load(f)
            nonpulsar_feature = dict(nonpulsar_feature, **temp_dict)
    print(len(nonpulsar_feature))
    for k, v in nonpulsar_feature.items():
        # print(v.shape)
        nonpulsar_feature[k] = v[:N, :M].reshape((1, N, M))

    # split data sets into two parts: train and test
    neg_index = nonpulsar_feature.keys()
    train_neg_index = neg_index[:TRAIN_NEG_NUM]
    test_neg_index = neg_index[TRAIN_NEG_NUM:]

    # put data into lmdb
    # lmdb_path = '/home/sunzy/workspace/data/MedlatTrainingData/lmdb/'
    lmdb_path = '/usr/data/MedlatTrainingData/lmdb/'
    # the train data db
    train_lmdb_path = os.path.join(lmdb_path, "%s_train_lmdb" % name)
    k = len(train_neg_index)
    put_into_lmdb(pulsar_feature, nonpulsar_feature, train_pos_index, train_neg_index, train_lmdb_path, k)
    # the test data db
    test_lmdb_path = os.path.join(lmdb_path, "%s_test_lmdb" % name)
    put_into_lmdb(pulsar_feature, nonpulsar_feature, test_pos_index, test_neg_index, test_lmdb_path, len(test_neg_index))

def put_into_lmdb(pos_dataset, neg_dataset, pos_index, neg_index, lmdb_file_path, k):
    """

    :param pos_dataset: positive samples data set. dict.
    :param neg_dataset: negative samples data set. dict.
    :param pos_index: index of positive samples. list.
    :param neg_index: index of negative samples. list.
    :param lmdb_file_path: the file path of the lmdb
    :param k: the num of negative samples actually in the data sets
    :return:
    """
    import random
    npos = len(pos_index)
    n = len(pos_index) + k
    index = list(range(n))
    random.shuffle(index)
    print(index)
    x = np.zeros((n, 1, N, M), dtype=np.float64)
    map_size = x.nbytes * 5

    env = lmdb.open(lmdb_file_path, map_size=map_size)

    with env.begin(write=True) as txn:
        # txn is a Transaction object
        # shuffle and put
        for i in range(n):
            datum = caffe.proto.caffe_pb2.Datum()
            datum.channels = x.shape[1]
            datum.height = x.shape[2]
            datum.width = x.shape[3]
            if index[i] < npos:
                # positive sample
                # print("%d %d pos sample" % (i, index[i]))
                datum.data = pos_dataset[pos_index[index[i]]].tobytes()
                datum.label = 1
            else:
                # negative sample
                # print("%d %d neg sample" % (i, index[i]))
                datum.data = neg_dataset[neg_index[index[i]-npos]].tobytes()
                datum.label = 0
            str_id = '{:08}'.format(i)
            txn.put(str_id.encode('ascii'), datum.SerializeToString())


        # # put train pos samples
        # for i in range(len(pos_index)):
        #     print(lmdb_file_path+":p:%d" % i)
        #     datum = caffe.proto.caffe_pb2.Datum()
        #     datum.channels = x.shape[1]
        #     datum.height = x.shape[2]
        #     datum.width = x.shape[3]
        #     datum.data = pos_dataset[pos_index[i]].tobytes()
        #     # datum.data = x[i].tobytes()  # or .tostring() if numpy < 1.9
        #     datum.label = 1
        #     str_id = '{:08}'.format(i)
        #
        #     # The encode is only essential in Python 3
        #     txn.put(str_id.encode('ascii'), datum.SerializeToString())
        # n1 = len(pos_index)
        # for i in range(k):
        #     print(lmdb_file_path+":p:%d" % i)
        #     datum = caffe.proto.caffe_pb2.Datum()
        #     datum.channels = x.shape[1]
        #     datum.height = x.shape[2]
        #     datum.width = x.shape[3]
        #     datum.data = neg_dataset[neg_index[i]].tobytes()
        #     # datum.data = x[i].tobytes()  # or .tostring() if numpy < 1.9
        #     datum.label = 0
        #     str_id = '{:08}'.format(i+n1)
        #
        #     # The encode is only essential in Python 3
        #     txn.put(str_id.encode('ascii'), datum.SerializeToString())

    env.close()



def test_lmdb():
    # lmdb_path = '/home/sunzy/workspace/data/MedlatTrainingData/lmdb'
    lmdb_path = '/usr/data/MedlatTrainingData/lmdb'

    img_lmdb = lmdb.open(os.path.join(lmdb_path, "%s_train_lmdb" % name))
    txn = img_lmdb.begin()
    cursor = txn.cursor()
    datum = caffe.proto.caffe_pb2.Datum()
    num = 0
    num2 = 0
    for (idx, (key, value)) in enumerate(cursor):
        datum.ParseFromString(value)
        flat_x = np.fromstring(datum.data, dtype=np.float64)
        # print(key)
        x = flat_x.reshape(datum.channels, datum.height, datum.width)
        y = datum.label
        num += 1

        # print(idx, key, x.shape, y)
        if y == 1:
            num2 += 1
    print(num, num2)

def main():
    # load_into_lmdb()
    test_lmdb()


if __name__ == "__main__":
    main()